import cli
import json
import os
import numpy as np
import pendulum
from dict_helpers import has_key, set_value

def check_mandatory_params(config, section, *keys):
	for key in keys:
		if not key in config[section]:
			cli.error(f'Parameter {section}.{key} is not in configuration file!')

def change_to_array(config, keys):
	if not has_key(config, keys): return
	if len(keys) == 1:
		if type(config[keys[0]]) != list:
			config[keys[0]] = [config[keys[0]]]
	elif len(keys) == 2:
		if type(config[keys[0]][keys[1]]) != list:
			config[keys[0]][keys[1]] = [config[keys[0]][keys[1]]]
	elif len(keys) == 3:
		if type(config[keys[0]][keys[1]][keys[2]]) != list:
			config[keys[0]][keys[1]][keys[2]] = [config[keys[0]][keys[1]][keys[2]]]

def check_array_size(config, keys, size):
	if not has_key(config, keys): return
	if len(keys) == 1:
		if len(config[keys[0]]) != size:
			cli.error(f'Parameter max_dom is {config["domains"]["max_dom"]}, but {keys} only has {len(config[keys[0]])} elements!')
	elif len(keys) == 2:
		if len(config[keys[0]][keys[1]]) != size:
			cli.error(f'Parameter max_dom is {config["domains"]["max_dom"]}, but {keys} only has {len(config[keys[0]][keys[1]])} elements!')
	elif len(keys) == 3:
		if len(config[keys[0]][keys[1]][keys[2]]) != size:
			cli.error(f'Parameter max_dom is {config["domains"]["max_dom"]}, but {keys} only has {len(config[keys[0]][keys[1]][keys[2]])} elements!')

def parse_config(config_json):
	config = {}

	if os.path.isfile(config_json):
		config = json.loads(open(config_json, 'r').read())
	else:
		try:
			config = json.loads(config_json)
		except:
			cli.error(f'{cli.red(config_json)} is not a JSON file or text!')

	if not has_key(config, 'custom'): config['custom'] = {}
	check_mandatory_params(config, 'custom', 'start_time')
	config['custom']['start_time'] = pendulum.from_format(config['custom']['start_time'], 'YYYYMMDDHH')
	if has_key(config, ('custom', 'forecast_hours')):
		config['custom']['end_time'] = config['custom']['start_time'].add(hours=config['custom']['forecast_hours'])
	else:
		check_mandatory_params(config, 'custom', 'end_time')
		config['custom']['end_time'] = pendulum.from_format(config['custom']['end_time'], 'YYYYMMDDHH')

	# If user set custom/domains/d01/{minlon, maxlon, minlat, maxlat}, we should set a proper projection and grid size accordingly.
	if has_key(config['custom'], 'domains') and type(config['custom']['domains']) == dict and len(config['custom']['domains']) > 0:
		check_mandatory_params(config['custom']['domains'], 'd01')
		check_mandatory_params(config['custom']['domains']['d01'], 'minlon', 'maxlon', 'minlat', 'maxlat')
		minlon = config['custom']['domains']['d01']['minlon']
		maxlon = config['custom']['domains']['d01']['maxlon']
		minlat = config['custom']['domains']['d01']['minlat']
		maxlat = config['custom']['domains']['d01']['maxlat']
		dlat = maxlat - minlat
		# Assume Lambert projection.
		if not has_key(config, 'geogrid'): config['geogrid'] = {}
		config['geogrid']['map_proj' ] = 'lambert'
		config['geogrid']['ref_lon'  ] = (minlon + maxlon) / 2
		config['geogrid']['ref_lat'  ] = (minlat + maxlat) / 2
		config['geogrid']['stand_lon'] = (minlon + maxlon) / 2
		config['geogrid']['truelat1' ] = minlat + 0.2 * dlat
		config['geogrid']['truelat2' ] = maxlat - 0.2 * dlat
		# Assume 3km resolution.
		R = 6371220 # Earth radius
		if has_key(config['custom']['domains']['d01'], 'dx'):
			dx = config['custom']['domains']['d01']['dx']
		elif has_key(config, ('domains', 'dx')):
			dx = config['domains']['dx']
		else:
			cli.error('No parameter dx is given!')
		if has_key(config['custom']['domains']['d01'], 'dy'):
			dy = config['custom']['domains']['d01']['dy']
		elif has_key(config, ('domains', 'dy')):
			dy = config['domains']['dy']
		else:
			cli.error('No parameter dy is given!')
		xspan = (np.radians(maxlon) - np.radians(minlon)) * R * np.cos(np.radians(config['geogrid']['ref_lat']))
		yspan = (np.radians(maxlat) - np.radians(minlat)) * R
		config['domains']['dx'] = dx
		config['domains']['dy'] = dy
		config['domains']['e_we'] = int(xspan / dx)
		config['domains']['e_sn'] = int(yspan / dy)
		if not has_key(config, ('domains', 'e_vert')): config['domains']['e_vert'] = 45

	# Set defaults.
	if not has_key(config, 'share'): config['share'] = {}
	if not has_key(config, 'domains'): config['domains'] = {}
	check_mandatory_params(config, 'domains', 'max_dom', 'dx', 'dy', 'e_we', 'e_sn')
	# - parent_grid_ratio, parent_time_step_ratio, i_parent_start, j_parent_start
	if config['domains']['max_dom'] == 1:
		config['domains']['parent_grid_ratio'] = [1]
		config['domains']['parent_time_step_ratio'] = [1]
		config['domains']['i_parent_start'] = [1]
		config['domains']['j_parent_start'] = [1]
	else:
		if has_key(config, ('custom', 'start_time')):
			check_mandatory_params(config, 'domains', 'parent_time_step_ratio')
		check_mandatory_params(config, 'domains', 'i_parent_start', 'j_parent_start')
	# Change to array.
	change_to_array(config, ('domains', 'dx'                    ))
	change_to_array(config, ('domains', 'dy'                    ))
	change_to_array(config, ('domains', 'e_we'                  ))
	change_to_array(config, ('domains', 'e_sn'                  ))
	change_to_array(config, ('domains', 'e_vert'                ))
	change_to_array(config, ('domains', 'parent_time_step_ratio'))
	change_to_array(config, ('domains', 'i_parent_start'        ))
	change_to_array(config, ('domains', 'j_parent_start'        ))
	# Check dimension.
	check_array_size(config, ('domains', 'dx'                    ), config['domains']['max_dom'])
	check_array_size(config, ('domains', 'dy'                    ), config['domains']['max_dom'])
	check_array_size(config, ('domains', 'e_we'                  ), config['domains']['max_dom'])
	check_array_size(config, ('domains', 'e_sn'                  ), config['domains']['max_dom'])
	check_array_size(config, ('domains', 'e_vert'                ), config['domains']['max_dom'])
	check_array_size(config, ('domains', 'parent_time_step_ratio'), config['domains']['max_dom'])
	check_array_size(config, ('domains', 'i_parent_start'        ), config['domains']['max_dom'])
	check_array_size(config, ('domains', 'j_parent_start'        ), config['domains']['max_dom'])
	if not has_key(config, ('domains', 'parent_grid_ratio')):
		config['domains']['parent_grid_ratio'] = [1]
		for i in range(1, len(config['domains']['dx'])):
			config['domains']['parent_grid_ratio'].append(int(config['domains']['dx'][i-1] / config['domains']['dx'][i]))
	config['domains']['grid_id'] = [i + 1 for i in range(config['domains']['max_dom'])]
	if not has_key(config, ('domains', 'parent_id')):
		config['domains']['parent_id'] = [i for i in range(config['domains']['max_dom'])]
		config['domains']['parent_id'][0] = 1

	# Transform parameters.
	if not type(config['domains']['dx']) == list: config['domains']['dx'] = [config['domains']['dx']]
	if not type(config['domains']['dy']) == list: config['domains']['dy'] = [config['domains']['dy']]

	if not has_key(config, 'geogrid'): config['geogrid'] = {}
	# - Check projection parameters.
	if not has_key(config, ('geogrid', 'map_proj')):
		config['geogrid']['map_proj'] = 'lambert'
		cli.warning(f'Set domains.map_proj to {config["domains"]["map_proj"]}.')
	if config['geogrid']['map_proj'] == 'lambert':
		check_mandatory_params(config, 'geogrid', 'ref_lon', 'ref_lat', 'truelat1', 'truelat2', 'stand_lon')
	for key in ('e_we', 'e_sn', 'parent_id', 'parent_grid_ratio', 'i_parent_start', 'j_parent_start'):
		config['geogrid'][key] = config['domains'][key]
	config['geogrid']['dx'] = config['domains']['dx'][0]
	config['geogrid']['dy'] = config['domains']['dy'][0]

	# Set time step if not set yet.
	if not has_key(config, ('domains', 'time_step')):
		if config['domains']['dx'][0] >= 30000:
 			config['domains']['time_step'] = 120
		elif 25000 <= config['domains']['dx'][0] < 30000:
			config['domains']['time_step'] = 120
		elif 20000 <= config['domains']['dx'][0] < 25000:
			config['domains']['time_step'] = 120
		elif 15000 <= config['domains']['dx'][0] < 20000:
			config['domains']['time_step'] = 90
		elif 10000 <= config['domains']['dx'][0] < 15000:
			config['domains']['time_step'] = 60
		elif 5000 <= config['domains']['dx'][0] < 10000:
			config['domains']['time_step'] = 30
		elif 2500 <= config['domains']['dx'][0] >= 5000:
			config['domains']['time_step'] = 10
		else:
			config['domains']['time_step'] = 5

	if not has_key(config, 'time_control'): config['time_control'] = {}
	if not 'history_interval' in config['time_control']:
		config['time_control']['history_interval'] = [60 for _ in range(config['domains']['max_dom'])]
	change_to_array(config, ('time_control', 'history_interval'))

	# wrfvar* sections
	if not has_key(config, 'wrfvar7'):
		config['wrfvar7'] = {
			'cv_options': 3
		}
	for tag in range(1, 23):
		section = f'wrfvar{tag}'
		if not section in config: config[section] = {}
	if not has_key(config, ('wrfvar3', 'ob_format')):
		set_value(config, ['wrfvar3', 'ob_format'], 2)
	if config['wrfvar3']['ob_format'] == 1 and not has_key(config, ('custom', 'wrfda', 'prepbufr_source')):
		set_value(config, ['custom', 'wrfda', 'prepbufr_source'], 'gdas')

	return config
